# -*- coding: utf-8 -*-
"""
会话数据服务 - 从数据库获取和管理水文模型会话数据

@author: AI Assistant
"""
import json
from typing import Optional, Dict, Any
from loguru import logger

# 导入配置和模型
import sys
from pathlib import Path
project_root = Path(__file__).resolve().parent.parent
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))

from utils.postgres_util import get_postgres_data
from models.domain_models import HydroModelSession
from utils.schema_parser import SchemaParser


class SessionService:
    """会话数据服务类"""
    
    def __init__(self):
        """初始化会话服务"""
        pass
    
    def get_session_by_id(self, session_id: str) -> Optional[HydroModelSession]:
        """从 t_hydromodel_session 表获取会话数据
        
        Args:
            session_id: 会话ID
            
        Returns:
            HydroModelSession 对象，如果未找到则返回 None
        """
        try:
            # 获取数据库连接
            conn = get_postgres_data()
            cursor = conn.cursor()
            
            # 查询SQL
            sql = """
                SELECT 
                    session_id, stcd, stcd_name, model_type, plcd,
                    initial_state, parameters, para_schema, createdate, md5
                FROM t_hydromodel_session
                WHERE session_id = %s
            """
            
            cursor.execute(sql, (session_id,))
            result = cursor.fetchone()
            
            if not result:
                logger.warning(f"未找到会话: {session_id}")
                cursor.close()
                conn.close()
                return None
            
            # 解析结果
            columns = [desc[0] for desc in cursor.description]
            session_data = dict(zip(columns, result))
            
            # 处理 JSONB 字段
            for field in ['initial_state', 'parameters', 'para_schema']:
                if session_data.get(field):
                    if isinstance(session_data[field], str):
                        try:
                            session_data[field] = json.loads(session_data[field])
                        except json.JSONDecodeError:
                            logger.warning(f"解析 {field} 字段失败，使用空字典")
                            session_data[field] = {}
                else:
                    session_data[field] = {}
            
            # 如果 para_schema 是简化格式（包含 parameters/initialState 作为字符串描述），转换为结构化格式
            para_schema = session_data.get('para_schema', {})
            if para_schema and 'parameters' in para_schema:
                # 检查是否是简化格式（值为字符串）
                first_param = list(para_schema['parameters'].values())[0] if para_schema['parameters'] else None
                if isinstance(first_param, str):
                    logger.info("检测到简化格式的 para_schema，正在转换为结构化格式...")
                    session_data['para_schema'] = SchemaParser.convert_simple_schema_to_structured(para_schema)
            
            # 创建 HydroModelSession 对象
            session = HydroModelSession(**session_data)
            
            logger.info(f"成功加载会话: {session_id}, 模型类型: {session.model_type}, 站点: {session.stcd}")
            
            cursor.close()
            conn.close()
            
            return session
            
        except Exception as e:
            logger.error(f"获取会话数据失败: {str(e)}", exc_info=True)
            if 'conn' in locals():
                conn.close()
            return None
    
    def update_session_parameters(
        self, 
        session_id: str, 
        initial_state: Optional[Dict[str, Any]] = None,
        parameters: Optional[Dict[str, Any]] = None
    ) -> bool:
        """更新会话参数（可选功能）
        
        Args:
            session_id: 会话ID
            initial_state: 新的初始状态（如果提供）
            parameters: 新的模型参数（如果提供）
            
        Returns:
            bool: 更新成功返回 True
        """
        try:
            if not initial_state and not parameters:
                logger.warning("没有提供要更新的参数")
                return False
            
            conn = get_postgres_data()
            cursor = conn.cursor()
            
            # 构建更新SQL
            update_fields = []
            update_values = []
            
            if initial_state is not None:
                update_fields.append("initial_state = %s")
                update_values.append(json.dumps(initial_state, ensure_ascii=False))
            
            if parameters is not None:
                update_fields.append("parameters = %s")
                update_values.append(json.dumps(parameters, ensure_ascii=False))
            
            if not update_fields:
                return False
            
            # 添加 session_id 到参数列表
            update_values.append(session_id)
            
            sql = f"""
                UPDATE t_hydromodel_session
                SET {', '.join(update_fields)}
                WHERE session_id = %s
            """
            
            cursor.execute(sql, tuple(update_values))
            conn.commit()
            
            affected_rows = cursor.rowcount
            cursor.close()
            conn.close()
            
            if affected_rows > 0:
                logger.info(f"成功更新会话参数: {session_id}")
                return True
            else:
                logger.warning(f"会话不存在或无更新: {session_id}")
                return False
                
        except Exception as e:
            logger.error(f"更新会话参数失败: {str(e)}", exc_info=True)
            if 'conn' in locals():
                conn.rollback()
                conn.close()
            return False
    
    def get_session_dict(self, session_id: str) -> Optional[Dict[str, Any]]:
        """获取会话数据（字典格式）
        
        Args:
            session_id: 会话ID
            
        Returns:
            会话数据字典，如果未找到则返回 None
        """
        session = self.get_session_by_id(session_id)
        if session:
            return session.model_dump()
        return None


# 使用示例
if __name__ == "__main__":
    service = SessionService()
    
    # 测试获取会话
    test_session_id = "test-session-123"
    session = service.get_session_by_id(test_session_id)
    
    if session:
        print(f"✅ 成功加载会话:")
        print(f"  - Session ID: {session.session_id}")
        print(f"  - 模型类型: {session.model_type}")
        print(f"  - 站点: {session.stcd} ({session.stcd_name})")
        print(f"  - 初始状态: {session.initial_state}")
        print(f"  - 参数: {session.parameters}")
    else:
        print(f"❌ 未找到会话: {test_session_id}")

